import pdfplumber
import PyPDF2
from io import BytesIO
from typing import List, Dict, Optional, Union
from .base_parser import BaseDocumentParser
import logging
from pathlib import Path

class PDFParser(BaseDocumentParser):
    """PDF文档解析器
    
    用于解析PDF文件中的文本内容和图片。支持多种图片格式的提取，包括JPG、PNG和JP2。
    
    Attributes:
        supported_image_formats (set): 支持的图片格式集合
    """
    
    def __init__(self, *args, **kwargs):
        """初始化PDF解析器
        
        Args:
            *args: 传递给父类的位置参数
            **kwargs: 传递给父类的关键字参数
        """
        super().__init__(*args, **kwargs)
        self.supported_image_formats = {'jpg', 'png', 'jp2'}
        self.logger = logging.getLogger(__name__)

    def extract_text(self, file_path: str) -> str:
        """提取PDF文本内容
        
        Args:
            file_path (str): PDF文件路径
            
        Returns:
            str: 提取的文本内容，包含图片占位符
            
        Raises:
            FileNotFoundError: 当文件不存在时
            PermissionError: 当没有文件读取权限时
        """
        if not Path(file_path).exists():
            raise FileNotFoundError(f"PDF文件不存在: {file_path}")
            
        text_content = []
        image_count = {}  # 记录每页的图片数量
        
        try:
            with pdfplumber.open(file_path) as pdf:
                for page_num, page in enumerate(pdf.pages, 1):
                    # 提取文本
                    text = page.extract_text()
                    if text:
                        text_content.append(text.strip())
                    
                    # 添加图片占位符
                    image_count[page_num] = 0
                    for image in page.images:
                        image_count[page_num] += 1
                        placeholder = f"[图片_{page_num}_{image_count[page_num]}]"
                        text_content.append(placeholder)
        
            return "\n".join(text_content)
        except Exception as e:
            self.logger.error(f"提取PDF文本时出错: {str(e)}")
            raise
    
    def extract_images(self, file_path: str) -> List[Dict[str, Union[bytes, str, int]]]:
        """提取PDF中的图片
        
        Args:
            file_path (str): PDF文件路径
            
        Returns:
            List[Dict]: 包含图片信息的字典列表，每个字典包含:
                - image_data (bytes): 图片二进制数据
                - format (str): 图片格式 ('jpg', 'png', 'jp2')
                - page (int): 页码
                - index (int): 图片索引
                
        Raises:
            FileNotFoundError: 当文件不存在时
            PermissionError: 当没有文件读取权限时
        """
        if not Path(file_path).exists():
            raise FileNotFoundError(f"PDF文件不存在: {file_path}")
            
        images = []
        
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                
                for page_num, page in enumerate(pdf_reader.pages, 1):
                    try:
                        if '/Resources' in page and '/XObject' in page['/Resources']:
                            xObject = page['/Resources']['/XObject'].get_object()
                            
                            for obj in xObject:
                                try:
                                    if xObject[obj]['/Subtype'] == '/Image':
                                        image = xObject[obj]
                                        
                                        # 获取图片数据和格式
                                        img_data = None
                                        img_format = None
                                        
                                        # 检查图片过滤器
                                        if '/Filter' in image:
                                            filters = image['/Filter']
                                            # 处理多个过滤器的情况
                                            if isinstance(filters, list):
                                                filters = filters[0]
                                                
                                            if filters == '/DCTDecode':
                                                img_data = image._data
                                                img_format = 'jpg'
                                            elif filters == '/FlateDecode':
                                                img_data = image._data
                                                img_format = 'png'
                                            elif filters == '/JPXDecode':
                                                img_data = image._data
                                                img_format = 'jp2'
                                        else:
                                            # 如果没有过滤器，尝试作为原始图像处理
                                            img_data = image._data
                                            img_format = 'png'
                                        
                                        # 只有当成功获取图片数据和格式时才添加
                                        if img_data and img_format:
                                            images.append({
                                                'image_data': img_data,
                                                'format': img_format,
                                                'page': page_num,
                                                'index': len(images) + 1
                                            })
                                except Exception as e:
                                    # 跳过处理有问题的图片
                                    continue
                    except Exception as e:
                        # 跳过处理有问题的页面
                        continue
            
            return images
        except Exception as e:
            self.logger.error(f"提取PDF图片时出错: {str(e)}")
            raise

    def _process_image_object(self, image: object, page_num: int, image_index: int) -> Optional[Dict]:
        """处理单个图片对象
        
        Args:
            image: PDF图片对象
            page_num (int): 页码
            image_index (int): 图片索引
            
        Returns:
            Optional[Dict]: 处理成功返回图片信息字典，失败返回None
        """
        try:
            if '/Filter' in image:
                filters = image['/Filter']
                if isinstance(filters, list):
                    filters = filters[0]
                
                format_map = {
                    '/DCTDecode': 'jpg',
                    '/FlateDecode': 'png',
                    '/JPXDecode': 'jp2'
                }
                
                img_format = format_map.get(filters, 'png')
                
                return {
                    'image_data': image._data,
                    'format': img_format,
                    'page': page_num,
                    'index': image_index
                }
            
            return None
        except Exception as e:
            print(f"处理图片时出错: {str(e)}")
            return None 